"""
Description: A simple teleoperation example using Vive Ultimate tracker to control a robot in simulation.

Author: Jin Liu
Date: 2025/06/18
"""
from pathlib import Path
import numpy as np

from ezgl import tb, Matrix4x4
from ezrobot.sim import RobotClient
from ezrobot.utils import GraphRunner, HandEyeConfig
from ezrobot.motion import OnlineMotion
from ezrobot.sensor import ViveTracker, SensorManager
from ezrobot.robot import UR5

BASE_DIR = Path(__file__).resolve().parent
np.set_printoptions(precision=3, suppress=True)


if __name__ == '__main__':
    # real robot
    robot = UR5(arm_ip="192.168.77.36", with_gripper=False)
    init_pose = robot.get_pose()

    # sim robot and motion
    client = RobotClient()
    client.clear_items()
    client.toggle_traj(True)

    # vive tracker
    vive = SensorManager(ViveTracker(device_index=1), frequency=30, buffer_size=10, streaming=True)
    handeye_config = HandEyeConfig.load(BASE_DIR / "../configs/handeye_config.yaml")
    
    # visualize robot
    visual_robot_runner = (
        GraphRunner()
        .add_node(0, robot.get_joints_rad)
        .add_node(1, client.set_joint, -1)
        .run_async(30)
    )

    # 遥操作
    def get_target_pose(vive2hand: Matrix4x4 = Matrix4x4.fromTranslation(-0.75, 0.30, 0)):
        """
        获取目标位姿, vive2hand
        """
        return vive2hand * handeye_config.base2map * vive.get_data() * handeye_config.hand2eye.inverse()
    
    start = []
    def move(target_pose):
        print(init_pose.toVector7d()-target_pose.toVector7d())
        if not start:
            if np.allclose(init_pose.toVector7d(), target_pose.toVector7d(), rtol=0, atol=0.15):
                start.append(True)
                print("Start teleoperation!")
        else:
            robot.servoL(target_pose, vel=0.1)

    vive_teleop_runner = (
        GraphRunner()
        .add_node(0, get_target_pose)
        .add_node(1, move, -1)
        .add_node(2, lambda p: client.set_pose_axis(p * Matrix4x4.fromScale(4, 4, 4)), 0)
    )

    vive_teleop_runner.run(50)